#-*- coding: utf-8 -*-
#
# Copyright 2013-2014 Antoine Drouin (poinix@gmail.com)
#
# This file is part of PAT.
#
#    PAT is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    PAT is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with PAT.  If not, see <http://www.gnu.org/licenses/>.
#

"""
 This is a 3dof (longitudinal) model for a fixed wing vehicle
"""
import math
import numpy as np
import scipy.optimize
import matplotlib.pyplot as plt

import pat.utils         as paut
import pat.frames        as pafr
import pat.atmosphere    as patm
import pat.dynamic_model as padm

"""
Components of the state vector
"""
sv_x     = 0 # position x axis in meter
sv_z     = 1 # position z axis in meter
sv_va    = 2 # air velocity in m/s
sv_gamma = 3 # air flight path ange in radians
sv_alpha = 4 # angle of attack in radians
sv_q     = 5 # angular velocity in rad/s
sv_size  = 6 # dimension of the state vector

"""
Components of the input vector
"""
iv_th   = 0 # position of throttle in percent
iv_de   = 1 # deflection of the elevator in rad
iv_size = 2 # dimension of the input vector


def aero_to_body(X):
    """ returns rotation matrix from aero to body"""
    ca = math.cos(X[sv_alpha]); sa = math.sin(X[sv_alpha])
    return np.array([[ca, sa],[-sa, ca]])

def aero_to_inert(X):
    """ returns rotation matrix from aero to inertial"""
    cg = math.cos(X[sv_gamma]); sg = math.sin(X[sv_gamma])
    return np.array([[cg, sg],[-sg, cg]])
 
def inert_to_body(X):
    """ returns rotation matrix from inertial to body"""
    theta = X[sv_gamma] + X[sv_alpha]
    ct = math.cos(theta); st = math.sin(theta)
    return np.array([[ct, -st],[st, ct]])

def body_to_inert(X):
    """ returns rotation matrix from body to inertial"""
    return np.transpose(inert_to_body(X))

def get_forces_inert(X, U, P, rho, Pdyn):
    """return Aero, Propulsion and Weight forces in body frame"""
    CL = P.CL0 +                             \
         P.CLalpha*(X[sv_alpha]-P.alpha0) +  \
         P.CLq*X[sv_q] +                     \
         P.CLde*U[iv_de]
    CD = P.CD0 + P.CDk1*CL + P.CDk2*(CL**2) +\
         P.CDde*U[iv_de] 
    T  = U[iv_th] * P.fmax *           \
         math.pow((rho/P.rhoi),P.nrho)*\
         math.pow((X[sv_va]/P.Vi),P.nV)
    return [Pdyn*P.Sref*np.dot(aero_to_inert(X), [-CD, -CL]),
            np.dot(body_to_inert(X), [T, 0]),
            [0, P.m*P.g]]

def get_momentum(X, U, P, Pdyn):
    """return momentum"""
    Cm = P.Cm0 +\
         P.Cmalpha * (X[sv_alpha]-P.alpha0) +\
         P.Cmq     * X[sv_q] +\
         P.Cmde    * U[iv_de]
    return Pdyn*P.Sref*Cm*P.Cref

    
def dyn(X, t, U, P):
    """
    Dynamic model
    """
    rho  = patm.get_rho(-X[sv_z])
    Pdyn = 0.5*rho*X[sv_va]**2

    [Faero, Fprop, Fweight] = get_forces_inert(X, U, P, rho, Pdyn)
    M = get_momentum(X, U, P, Pdyn)
    # airspeed in earth frame
    Vae =  np.dot(aero_to_inert(X), [X[sv_va], 0])
    # time derivative of airspeed in earth frame
    Vdae = 1./P.m*(Faero+Fprop+Fweight) - P.Wde

    Xd = np.zeros(sv_size)
    Xd[sv_x:sv_z+1] = Vae + P.We
    Xd[sv_va] = (Vae[0]*Vdae[0] + Vae[1]*Vdae[1]) / X[sv_va]
    Xd[sv_gamma] = -(Vae[0]*Vdae[1] - Vae[1]*Vdae[0]) / X[sv_va]**2
    Xd[sv_alpha] = X[sv_q] - Xd[sv_gamma]
    Xd[sv_q] = 1./P.I*M
    return Xd


def trim(P, args=None):
    """
    Find throttle, elevator  and angle of attack corresponding
    to the given airspeed and and flight path
    """
    if args<>None:
        va, gamma, h = (args['va'], args['gamma'], args['h'] )
    else:
        va, gamma, h = (70., 0., 0.)

    def err_func((throttle, elevator, alpha)):
        X = [0., -h, va, gamma, alpha, 0.]; U = [throttle, elevator]
        Xdot = dyn(X, 0., U, P) 
        Xdot_ref = [va*math.cos(gamma), -va*math.sin(gamma), 0, 0, 0, 0]
        return np.linalg.norm(Xdot - Xdot_ref)
    trim0 = [0.2, paut.rad_of_deg(0.), paut.rad_of_deg(0.)]
    thr_e, ele_e, alpha_e = scipy.optimize.fmin_powell(err_func, trim0, disp=False)
    Xe = np.array([0, -h, va, gamma, alpha_e, 0])
    Ue = np.array([thr_e, ele_e])
    return Xe, Ue

def saturate_input(P, U):
    sats = np.array([[0, 1.],[paut.rad_of_deg(-30), paut.rad_of_deg(30)]])
    return paut.saturate(U, sats)

class Param:
    """
    Parameters of the dynamic model
    """
    
    def __init__(self):
        self.name    = 'Default'
        self.g       = 9.81        # gravitiy in m/s2
        # Wind
        self.We      = np.zeros(2) # wind in earth frame in m/s
        self.Wde     = np.zeros(2) # wind time derivative
        # Inertia
        self.m       = 1000.       # mass in kg
        self.I       = 3000.       # inertia in 
        # Aerodynamics
        self.Sref    = 15.         # reference surface in m2
        self.Cref    = 1.          # reference chord in m
        self.alpha0  = 0.          # alpha0
        self.CL0     = 0.09        # CL0
        self.CLalpha = 4.95        # CLalpha
        self.CLq     = 9.2         # CLq
        self.CLde    = 0.48        # CLde
        self.CD0     = 0.02        # CD0
        self.CDk1    =-0.00017     # CDk1
        self.CDk2    =-0.05        # CDk2
        self.CDde    = 0           # CDde
        self.Cm0     = 0           # Cm0
        self.Cmalpha = -2.46       # Cmalpha
        self.Cmq     = -10.5       # Cmq
        self.Cmde    = -1.23       # Cmde
        # Propulsion
        self.fmax    = 6000        # max thrust of engine in N
        self.rhoi    = 1.          # reference density
        self.nrho    = 0           # 
        self.Vi      = 1.          # reference velocity
        self.nV      = 0.          #
        self.eng_to_body = np.array([[1., 0], [0., 1.]])



class DynamicModel(padm.DynamicModel):

    sv_x     = sv_x
    sv_z     = sv_z
    sv_va    = sv_va
    sv_gamma = sv_gamma
    sv_alpha = sv_alpha
    sv_q     = sv_q
    sv_size  = sv_size

    iv_th    = iv_th
    iv_de    = iv_de
    iv_size  = iv_size

    dyn  = lambda self, X, t, U, P: dyn(X, t, U, self.P)
    trim = lambda self, args=None: trim(self.P, args)
    plot_trajectory = lambda self, time, X, U=None: plot_trajectory(time, X, U)

    def __init__(self):
        self.P = Param()
        self.X = np.zeros(DynamicModel.sv_size)
        self.Xd = np.zeros(DynamicModel.sv_size)

    def state_SixDOFfEuclidianEuler(self):
        X = np.zeros(pafr.SixDOFfEuclidianEuler.size)
        X[pafr.SixDOFfEuclidianEuler.x]     =  self.X[sv_x]
        X[pafr.SixDOFfEuclidianEuler.z]     =  self.X[sv_z]
        X[pafr.SixDOFfEuclidianEuler.xd]    =  self.X[sv_va]*math.cos(self.X[sv_gamma])+self.P.We[0]
        X[pafr.SixDOFfEuclidianEuler.zd]    = -self.X[sv_va]*math.sin(self.X[sv_gamma])+self.P.We[1]
        X[pafr.SixDOFfEuclidianEuler.theta] =  self.X[sv_alpha] + self.X[sv_gamma]         
        X[pafr.SixDOFfEuclidianEuler.q]     =  self.X[sv_q]
        return X

    def name(self): return "Fixed Wing Longitudinal"
    

def plot_trajectory(time, X, U=None, figure=None, window_title="Trajectory",
                    legend=None, filename=None):
    plots = [("$x$",       "m",     X[:,sv_x]),
             ("$z$",       "m",     X[:,sv_z]),
             ("$v_a$",     "m/s",   X[:,sv_va]),
             ("$\\gamma$", "deg",   paut.deg_of_rad(X[:,sv_gamma])),
             ("$\\alpha$", "deg",   paut.deg_of_rad(X[:,sv_alpha])),
             ("$q$",       "deg/s", paut.deg_of_rad(X[:,sv_q]))]
    if U<>None:
        plots += [("$d_{th}$", "%",  100*U[:, iv_th]),
                  ("$d_e$",        "deg", paut.deg_of_rad(U[:, iv_de]))]
    
    return paut.plot_in_grid(time, plots, 2, figure, window_title, legend, filename,
                             margins=(0.04, 0.05, 0.98, 0.96, 0.10, 0.34))


def plot_trajectory_alt(time, X, U=None, figure=None, window_title="Trajectory",
                        legend=None, filename=None):
    plots = [("$x$",       "m",     X[:,sv_x]),
             ("$z$",       "m",     X[:,sv_z]),
             ("$v_a$",     "m/s",   X[:,sv_va]),
             ("$\\gamma$", "deg",   paut.deg_of_rad(X[:,sv_gamma])),
             ("$\\theta$", "deg",   paut.deg_of_rad(X[:,sv_alpha]+X[:,sv_gamma])),
             ("$q$",       "deg/s", paut.deg_of_rad(X[:,sv_q]))]
    if U<>None:
        plots += [("$d_{th}$", "%",  100*U[:, iv_th]),
                  ("$d_e$",        "deg", paut.deg_of_rad(U[:, iv_de]))]
        
    return paut.plot_in_grid(time, plots, 2, figure, window_title, legend, filename,
                             margins=(0.04, 0.08, 0.93, 0.92, 0.10, 0.34))


def plot_angular_dynamics_alt(time, X, U=None, figure=None, window_title="Trajectory",
                              legend=None, filename=None):
    plots = [("$\\theta$", "deg",   paut.deg_of_rad(X[:,sv_alpha]+X[:,sv_gamma])),
             ("$q$",       "deg/s", paut.deg_of_rad(X[:,sv_q]))]
    if U<>None:
        plots += [("$d_{th}$", "%",  100*U[:, iv_th]),
                  ("$d_e$",        "deg", paut.deg_of_rad(U[:, iv_de]))]
    return paut.plot_in_grid(time, plots, 2, figure, window_title, legend, filename,
                             margins=(0.04, 0.08, 0.99, 0.92, 0.10, 0.34))
    


